//      dist_main.cc
//      
//      Copyright 2012 tqlong <tqlong@espada>
//      
//      This program is free software; you can redistribute it and/or modify
//      it under the terms of the GNU General Public License as published by
//      the Free Software Foundation; either version 2 of the License, or
//      (at your option) any later version.
//      
//      This program is distributed in the hope that it will be useful,
//      but WITHOUT ANY WARRANTY; without even the implied warranty of
//      MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
//      GNU General Public License for more details.
//      
//      You should have received a copy of the GNU General Public License
//      along with this program; if not, write to the Free Software
//      Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
//      MA 02110-1301, USA.

/* example run
generate random A=W*H: 
   ./distnmf --rand --m=10 --k=2 --n=5 --output=d1 --verbose --binary
read matrix
   ./distnmf --read --input=d1 --binary
 */

//#include <iostream>
#include <boost/program_options.hpp>
#include <boost/date_time/posix_time/posix_time.hpp> 
#include <boost/foreach.hpp>

#include "helper/stl_helper.h"
#include "helper/ublas_helper.h"
#include "helper/ml_data.h"
#include "helper/rand_helper.h"

namespace po = boost::program_options;
po::options_description desc("Allowed options");
po::variables_map vm;

void process_options(int argc, char** argv);
void usage(int exitcode = 0);

void read_data();
void rand_data();
template<typename MATRIX> void nmf_bcd();

boost::posix_time::ptime time_start_all(boost::posix_time::microsec_clock::local_time());
boost::posix_time::ptime time_start(boost::posix_time::microsec_clock::local_time());

boost::posix_time::ptime now() 
{
  return boost::posix_time::microsec_clock::local_time();
}

boost::posix_time::time_duration duration() 
{
  boost::posix_time::time_duration d = now() - time_start;
  time_start = now();
  return d;
}

boost::posix_time::time_duration duration_all() 
{
  return now() - time_start_all;
}

int main(int argc, char** argv)
{
  //MSG("BOOST SUB-VERSION " << BOOST_VERSION / 100 % 100);

  process_options(argc, argv);

  if (vm.count("help"))
    usage();
  else if (vm.count("read"))
    read_data();
  else if (vm.count("rand"))
    rand_data();
  else if (vm.count("bcd"))
    nmf_bcd<matrix_t>();
  else
    usage();

  MSG("Duration = " << duration_all());
}

void process_options(int argc, char** argv)
{
  desc.add_options()
    ("help", "produce this usage")
    ("input", po::value<std::string>(), "input file")
    ("output", po::value<std::string>(), "output file")
    ("log", po::value<std::string>(), "log file")
    ("verbose", "print debug information")

    // RAND
    ("m", po::value<int>(), "number of rows of A")
    ("n", po::value<int>(), "number of cols of A")    
    ("k", po::value<int>(), "true dimension of W and H")
    ("s", po::value<double>()->default_value(0.5), "sparsity level of W and H")
    ("binary", "save output in binary format")

    // NMF
    ("alpha", po::value<double>(), "regularization parameter \alpha*|W|_F^2 + \alpha*|H|_F^2")
    ("iter", po::value<int>(), "maximum number of iterations")

    //JOB TYPE
    ("read", "read data file\n--input [--output]")
    ("rand", "random data file A=WH (size m x k x n)\n--m --n --k --output [--binary]")
    ("bcd", "Block Coordinate Descent\n--input [--binary] --k --alpha --output --log")
    ;
  po::store(po::parse_command_line(argc, argv, desc), vm);
  po::notify(vm);
}

void usage(int exitcode)
{
  std::cout << desc << std::endl;
  exit(exitcode);
}

template<typename MATRIX>
void gen_sparse_matrix(MATRIX& m, int rows, int cols, double sparsity, 
		       bool reset_seed = false) {
  m.resize(rows, cols, false);
  RAND::randu_t r = RAND::randu();
  if (reset_seed) {
    r.engine().seed(time(NULL));
    r.distribution().reset();
  }
  for (int i = 0; i < rows; i++)
    for (int j = 0; j < cols; j++)
      m(i,j) = (r() > sparsity) ? r() : 0;      
}

struct line {
  int i_;
  int j_;
  double s_;
  
  line() {}
  line(int i, int j, double s) : i_(i), j_(j), s_(s) {}      
  static void write_block(std::ostream& o, line l) {
    o.write((char*)&l, sizeof(line));
  }
  std::istream& read(std::istream& i) {
    return i.read((char*)this, sizeof(line));
  }
};

template<typename MATRIX>
void read_sparse_matrix(const std::string& inFile, MATRIX& A, bool is_binary) 
{
  FileAccessRead(in, inFile);
  if (!is_binary) {
    int m, n;
    double check_sum;
    in >> m >> n >> check_sum;
    A.resize(m, n, false);
    double sum = 0;
    for (int i = 0; i < m; i++)
      for (int j = 0; j < n; j++) {
	double s;
	in >> s;
	if ( fabs(s) > 1e-16 ) A(i,j) = s;
	sum += s;
      }
    MSG("check_sum=" << check_sum << " sum=" << sum << " diff=" << fabs(sum-check_sum));
    //if (fabs(sum-check_sum) > 1e-10) FATAL_ERROR("check sum error");
  }
  else { // binary mode
    line head, l;
    head.read(in);
    A.resize(head.i_, head.j_, false);
    double sum = 0;
    while (!in.eof()) {
      l.read(in);
      if (in.fail()) break;
      A(l.i_, l.j_) = l.s_;
      sum += l.s_;
    }
    MSG("check_sum=" << head.s_ << " sum=" << sum << " diff=" << fabs(sum-head.s_));
    //if (fabs(sum-head.s_) > 1e-10) FATAL_ERROR("check sum error");
  }
  in.close();
}

template<typename MATRIX>
void save_sparse_matrix(const std::string& outFile, const MATRIX& A, bool is_binary) {
  double sum = 0;
  tr1(A, i1)
    tr(i1, i2) sum += *i2;

  FileAccessWrite(out, outFile);
  if (!is_binary) { // text mode
    out << "\t" << A.size1() << "\t" << A.size2() << "\t" << sum << ENDL;
    for (uint i = 0; i < A.size1(); i++) {
      for (uint j = 0; j < A.size2(); j++) out << "\t" << A(i, j);
      out << ENDL;
    }
  }
  else { // binary mode
    //#define tr1(x, it) for ( typeof((x).begin1()) it = (x).begin1(); it != (x).end1(); it++)
    line::write_block(out, line(A.size1(), A.size2(), sum));
    tr1(A, i1) {
      tr(i1, i2)
	line::write_block(out, line(i2.index1(), i2.index2(), *i2));
    }
  }
  MSG("check sum=" << sum);
  out.close();
}

void read_data()
{
  std::string inFile = VM(vm, "input", std::string);
  sparse_matrix_t A;
  read_sparse_matrix<sparse_matrix_t>(inFile, A, vm.count("binary"));
}

void rand_data()
{
  std::string outFile = VM(vm, "output", std::string);
  int m = VM(vm, "m", int), 
    n = VM(vm, "n", int), 
    k = VM(vm, "k", int);
  double sparsity = VM(vm, "s", double);
  
  // generate random W and H
  sparse_matrix_t A, W, H;
  gen_sparse_matrix<sparse_matrix_t>(W, m, k, sparsity);
  gen_sparse_matrix<sparse_matrix_t>(H, k, n, sparsity);

  // A = W*H
  A = ublas::prod(W,H);
  double s_A = 0;
  tr1(A, i1) tr(i1, i2) s_A += *i2;
  s_A /= (A.size1() * A.size2());
  
  // add uniform random noise to A
  RAND::randn_t r = RAND::randn();
  RAND::randu_t r1 = RAND::randu();
  tr1(A, i1) tr(i1, i2) {
    double nr = *i2 + r() * (s_A/100);
    *i2 = (nr > 0) ? nr : 0;
  }

  MSG_VERBOSE(vm, "W=" << W << ENDL << "H=" << H << ENDL << "A=" << A);

  // save to output file
  save_sparse_matrix<sparse_matrix_t>(outFile, A, vm.count("binary"));
}

template<typename MATRIX>
MATRIX& nmf_bcd_W(const MATRIX& A, const MATRIX& W, const MATRIX& H, double alpha,
		  MATRIX& A_H, MATRIX& H_T_H, MATRIX& Wnew) {
  //MATRIX A_H(A.size1(), H.size1());
  ublas::axpy_prod(A, ublas::trans(H), A_H, true);

  //MATRIX H_T_H(H.size1(), H.size1());
  ublas::axpy_prod(H, ublas::trans(H), H_T_H, true);
  //Wnew = W;

  for (uint i = 0; i < W.size2(); i++) {
    double h_ii = H_T_H(i,i);
    vec_t w_i = col(A_H,i);
    ublas::axpy_prod(Wnew, -col(H_T_H,i), w_i, false);
    w_i /= (h_ii+alpha);
    w_i += col(Wnew,i)*h_ii/(h_ii+alpha);

    for (uint j = 0; j < W.size1(); j++)
      if (w_i(j) <= 1e-16) w_i(j) = 1e-16;
    col(Wnew,i) = w_i;
  }

  return Wnew;
}

template<typename MATRIX>
MATRIX& nmf_bcd_H(const MATRIX& A, const MATRIX& W, const MATRIX& H, double alpha,
		 MATRIX& A_W, MATRIX& W_T_W, MATRIX& Hnew) {
  //MATRIX A_W(A.size2(), W.size2());
  ublas::axpy_prod(ublas::trans(A), W, A_W, true);

  //MATRIX W_T_W(W.size2(), W.size2());
  ublas::axpy_prod(ublas::trans(W), W, W_T_W, true);
  //Hnew = H;

  for (uint i = 0; i < H.size1(); i++) {
    double w_ii = W_T_W(i,i);

    vec_t h_i = col(A_W,i);
    ublas::axpy_prod(ublas::trans(Hnew), -col(W_T_W,i), h_i, false);
    h_i /= (w_ii+alpha);
    h_i += ublas::trans(row(Hnew,i))*w_ii/(w_ii+alpha);

    for (uint j = 0; j < H.size2(); j++)
      if (h_i(j) <= 1e-16) h_i(j) = 1e-16;
    row(Hnew,i) = ublas::trans(h_i);
  }

  return Hnew;
}

template<typename MATRIX>
void nmf_bcd()
{
  std::string inFile = VM(vm, "input", std::string);
  std::string logFile = VM(vm, "log", std::string);
  MATRIX A;
  read_sparse_matrix<MATRIX>(inFile, A, vm.count("binary"));

  int m = A.size1(), n = A.size2(), k = VM(vm, "k", int);
  double alpha = VM(vm, "alpha", double);
  int maxiter = VM(vm, "iter", int);

  MATRIX W, H; 
  gen_sparse_matrix<MATRIX>(W, m, k, 0.5, true);
  gen_sparse_matrix<MATRIX>(H, k, n, 0.5, true);

  MATRIX A_H(A.size1(), H.size1());
  MATRIX A_W(A.size2(), W.size2());
  MATRIX W_T_W(W.size2(), W.size2());
  MATRIX H_T_H(H.size1(), H.size1());
  MATRIX Wnew = W, Hnew = H;
  
  FileAccessCreate(logFile);
  for (int iter = 0; iter < maxiter; iter++) {
    W = nmf_bcd_W<MATRIX>(A, W, H, alpha, A_H, H_T_H, Wnew);
    H = nmf_bcd_H<MATRIX>(A, W, H, alpha, A_W, W_T_W, Hnew);
    
    MATRIX Atmp(m, n);
    ublas::axpy_prod(W,H,Atmp,true);
    double err = ublas::norm_frobenius(A-Atmp);
    FileAccessAppend(log, logFile);
    log << "\t" << iter << "\t" << err << ENDL;
    if (iter % 10 == 0 || iter == maxiter-1)
      MSG("iter=" << iter << " err=" << err);
  }
}
